# three_smoke_modified.py
import numpy as np
import random
import math
import pandas as pd
import time

g = 9.8
missile_speed = 300.0
cloud_sink = 3.0
cloud_r = 10.0
cloud_duration = 20.0

# 定义所有导弹和无人机的位置
missile_positions = {
    'M1': np.array([20000.0, 0.0, 2000.0]),
    'M2': np.array([19000.0, 600.0, 2100.0]),
    'M3': np.array([18000.0, -600.0, 1900.0])
}

drone_positions = {
    'FY1': np.array([17800.0, 0.0, 1800.0]),
    'FY2': np.array([12000.0, 1400.0, 1400.0]),
    'FY3': np.array([6000.0, -3000.0, 700.0]),
    'FY4': np.array([11000.0, 2000.0, 1800.0]),
    'FY5': np.array([13000.0, -2000.0, 1300.0])
}

fake_target = np.array([0.0, 0.0, 0.0])
real_target = np.array([0.0, 200.0, 0.0])

# 全局变量，将在main函数中设置
alpha_fixed = None
v_fixed = None
current_missile = None
current_drone = None


# ------------------ 基本几何/物理函数 ------------------
def missile_pos(t, missile_name):
    M = missile_positions[missile_name]
    missile_dir = fake_target - M
    missile_unit = missile_dir / np.linalg.norm(missile_dir)
    return M + missile_unit * missile_speed * t


def explosion_center(alpha, v, td, f, drone_name):
    FY = drone_positions[drone_name]
    dir_xy = np.array([math.cos(alpha), math.sin(alpha)])
    drop = FY + np.array([v * dir_xy[0] * td, v * dir_xy[1] * td, 0.0])
    burst = drop + np.array([v * dir_xy[0] * f, v * dir_xy[1] * f, -0.5 * g * f * f])
    return burst


def cloud_center(alpha, v, td, f, tau, drone_name):
    burst = explosion_center(alpha, v, td, f, drone_name)
    z = max(0.0, burst[2] - cloud_sink * tau)
    return np.array([burst[0], burst[1], z])


def point_segment_dist(C, M, P):
    MP = P - M
    seg2 = MP.dot(MP)
    if seg2 <= 1e-12:
        return np.linalg.norm(C - M)
    t = np.dot(C - M, MP) / seg2
    t = max(0.0, min(1.0, t))
    proj = M + t * MP
    return np.linalg.norm(C - proj)


def blocked_at_time(alpha, v, td, f, t, drone_name, missile_name):
    te = td + f
    tau = t - te
    if tau < 0.0 or tau > cloud_duration:
        return False
    Mpos = missile_pos(t, missile_name)
    Cpos = cloud_center(alpha, v, td, f, tau, drone_name)
    d = point_segment_dist(Cpos, Mpos, real_target)
    return d <= cloud_r + 1e-9


# ------------------ 评价函数（硬判定） ------------------
def total_shield_time_from_plans(drops, bursts, drone_name, missile_name, dt=0.1):
    # 计算导弹飞行时间
    M = missile_positions[missile_name]
    missile_dir = fake_target - M
    missile_flight_time = np.linalg.norm(missile_dir) / missile_speed

    times = np.arange(0.0, missile_flight_time + 1e-9, dt)
    shield_flags = np.zeros_like(times, dtype=bool)
    indiv_flags = [np.zeros_like(times, dtype=bool) for _ in range(3)]

    for idx, t in enumerate(times):
        tglob = float(t)
        for i in range(3):
            tau = tglob - (drops[i] + bursts[i])
            if tau < 0 or tau > cloud_duration:
                continue
            if blocked_at_time(alpha_fixed, v_fixed, drops[i], bursts[i], tglob, drone_name, missile_name):
                indiv_flags[i][idx] = True
                shield_flags[idx] = True

    total = shield_flags.sum() * dt
    indiv_intervals = []
    indiv_times = []

    for arr in indiv_flags:
        intervals = []
        in_seg = False
        start = None
        for k, val in enumerate(arr):
            tcur = times[k]
            if val and not in_seg:
                in_seg = True
                start = tcur
            if not val and in_seg:
                intervals.append((start, times[k - 1] + dt))
                in_seg = False
                start = None
        if in_seg:
            intervals.append((start, times[-1] + dt))
        indiv_intervals.append(intervals)
        indiv_times.append(sum(e - s for s, e in intervals))

    return total, indiv_times, indiv_intervals


# ------------------ 单枚弹优化（SA） ------------------
def optimize_one_grenade(fixed_drops, fixed_bursts, idx_to_opt, drone_name, missile_name,
                         td_bounds=(0, 60), f_bounds=(0.5, 10.0),
                         dt_eval=0.1,
                         sa_steps=1500, T0=1.0, cooling=0.999):
    if idx_to_opt == 0:
        cur_td = 1.5
        cur_f = 3.6
    else:
        prev_td = max(0.0, fixed_drops[idx_to_opt - 1] if fixed_drops[idx_to_opt - 1] is not None else 2.0)
        cur_td = min(td_bounds[1], prev_td + 6.0)
        cur_f = 3.5

    cur_td = float(max(td_bounds[0], min(td_bounds[1], cur_td)))
    cur_f = float(max(f_bounds[0], min(f_bounds[1], cur_f)))

    drops = [fixed_drops[i] if fixed_drops[i] is not None else (cur_td if i == idx_to_opt else None) for i in range(3)]
    bursts = [fixed_bursts[i] if fixed_bursts[i] is not None else (cur_f if i == idx_to_opt else None) for i in
              range(3)]

    for i in range(3):
        if drops[i] is None:
            drops[i] = 0.0
        if bursts[i] is None:
            bursts[i] = 3.5

    best_td, best_f = cur_td, cur_f
    best_plan_drops = drops.copy()
    best_plan_bursts = bursts.copy()
    best_plan_drops[idx_to_opt] = best_td
    best_plan_bursts[idx_to_opt] = best_f

    best_score, _, _ = total_shield_time_from_plans(best_plan_drops, best_plan_bursts, drone_name, missile_name,
                                                    dt=dt_eval)

    T = T0
    cur_td, cur_f = best_td, best_f
    cur_score = best_score

    for step in range(sa_steps):
        if random.random() < 0.05:
            cand_td = random.uniform(*td_bounds)
            cand_f = random.uniform(*f_bounds)
        else:
            cand_td = cur_td + random.gauss(0, 1.0)
            cand_f = cur_f + random.gauss(0, 0.6)

        cand_td = max(td_bounds[0], min(td_bounds[1], cand_td))
        cand_f = max(f_bounds[0], min(f_bounds[1], cand_f))

        cand_drops = best_plan_drops.copy()
        cand_bursts = best_plan_bursts.copy()
        cand_drops[idx_to_opt] = cand_td
        cand_bursts[idx_to_opt] = cand_f

        ok = True
        if idx_to_opt > 0:
            prev_td = cand_drops[idx_to_opt - 1]
            if cand_td < prev_td + 1.0:
                ok = False

        # 计算导弹飞行时间
        M = missile_positions[missile_name]
        missile_dir = fake_target - M
        missile_flight_time = np.linalg.norm(missile_dir) / missile_speed

        if cand_td + cand_f > missile_flight_time - 0.01:
            ok = False

        if not ok:
            T *= cooling
            continue

        score, _, _ = total_shield_time_from_plans(cand_drops, cand_bursts, drone_name, missile_name, dt=dt_eval)

        if score > cur_score or random.random() < math.exp((score - cur_score) / max(1e-9, T)):
            cur_td, cur_f = cand_td, cand_f
            cur_score = score
            best_plan_drops[idx_to_opt] = cand_td
            best_plan_bursts[idx_to_opt] = cand_f
            if score > best_score:
                best_td, best_f = cand_td, cand_f
                best_score = score

        T *= cooling

    best_local_td, best_local_f, best_local_score = best_td, best_f, best_score

    # 在最优解附近进行网格搜索
    grid_td = np.arange(max(td_bounds[0], best_td - 2.0), min(td_bounds[1], best_td + 2.0) + 1e-9, 0.2)
    grid_f = np.arange(max(f_bounds[0], best_f - 1.5), min(f_bounds[1], best_f + 1.5) + 1e-9, 0.2)

    for tdg in grid_td:
        for fg in grid_f:
            cand_drops = best_plan_drops.copy()
            cand_bursts = best_plan_bursts.copy()
            cand_drops[idx_to_opt] = tdg
            cand_bursts[idx_to_opt] = fg

            if idx_to_opt > 0:
                if cand_drops[idx_to_opt] < cand_drops[idx_to_opt - 1] + 1.0:
                    continue

            # 计算导弹飞行时间
            M = missile_positions[missile_name]
            missile_dir = fake_target - M
            missile_flight_time = np.linalg.norm(missile_dir) / missile_speed

            if tdg + fg > missile_flight_time - 0.01:
                continue

            score, _, _ = total_shield_time_from_plans(cand_drops, cand_bursts, drone_name, missile_name, dt=dt_eval)

            if score > best_local_score:
                best_local_score = score
                best_local_td = tdg
                best_local_f = fg

    return best_local_td, best_local_f, best_local_score


# ------------------ 主：逐枚接力优化 ------------------
def sequential_optimize(drone_name, missile_name, save_excel=False):
    # 计算导弹飞行时间
    M = missile_positions[missile_name]
    missile_dir = fake_target - M
    missile_flight_time = np.linalg.norm(missile_dir) / missile_speed

    td_bounds = (0.0, min(60.0, missile_flight_time - 1.0))
    f_bounds = (0.5, 10.0)

    drops = [None, None, None]
    bursts = [None, None, None]

    print(f"优化{drone_name}对{missile_name}的第1枚烟幕弹...")
    td1, f1, s1 = optimize_one_grenade(drops, bursts, 0, drone_name, missile_name,
                                       td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=2000, T0=1.2, cooling=0.9995)
    drops[0] = td1
    bursts[0] = f1
    print(f"  弹1: td={td1:.2f}, f={f1:.2f}, 单独遮蔽(估)={s1:.2f}s")

    print(f"优化{drone_name}对{missile_name}的第2枚烟幕弹...")
    td2, f2, s2 = optimize_one_grenade(drops, bursts, 1, drone_name, missile_name,
                                       td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=1800, T0=1.0, cooling=0.9995)
    drops[1] = td2
    bursts[1] = f2
    print(f"  弹2: td={td2:.2f}, f={f2:.2f}, 累计遮蔽(估)={s2:.2f}s")

    print(f"优化{drone_name}对{missile_name}的第3枚烟幕弹...")
    td3, f3, s3 = optimize_one_grenade(drops, bursts, 2, drone_name, missile_name,
                                       td_bounds=td_bounds, f_bounds=f_bounds,
                                       dt_eval=0.1, sa_steps=1600, T0=1.0, cooling=0.9995)
    drops[2] = td3
    bursts[2] = f3
    print(f"  弹3: td={td3:.2f}, f={f3:.2f}, 累计遮蔽(估)={s3:.2f}s")

    total, indiv_times, indiv_intervals = total_shield_time_from_plans(drops, bursts, drone_name, missile_name, dt=0.02)

    print(f"\n=== {drone_name}对{missile_name}的最终精确评估 ===")
    print(f"飞行方向 (deg): {math.degrees(alpha_fixed):.2f}, 速度 (m/s): {v_fixed:.1f}")
    for i in range(3):
        print(
            f"弹{i + 1}: 投放={drops[i]:.2f}s, 引信={bursts[i]:.2f}s, 单独遮蔽={indiv_times[i]:.2f}s, 区间={indiv_intervals[i]}")
    print(f"并集总有效遮蔽时间: {total:.2f} s")

    return drops, bursts, total, indiv_times, indiv_intervals


def main(drone_name, missile_name, alpha_deg, v_value, save_excel=False):
    global alpha_fixed, v_fixed
    alpha_fixed = alpha_deg * math.pi / 180.0
    v_fixed = v_value

    random.seed(42)
    np.random.seed(42)

    drops, bursts, total, indiv_times, indiv_intervals = sequential_optimize(drone_name, missile_name,
                                                                             save_excel=save_excel)

    return total, drops, bursts, indiv_times


if __name__ == "__main__":
    # 示例：FY1对M1的情况
    drone_name = "FY1"
    missile_name = "M1"
    alpha_deg = -180
    v_value = 121.09

    random.seed(42)
    np.random.seed(42)

    t0 = time.time()
    total, drops, bursts, indiv_times = main(drone_name, missile_name, alpha_deg, v_value, save_excel=False)
    print("总耗时: {:.1f}s".format(time.time() - t0))